import gymnasium as gym
import numpy as np
from gymnasium import spaces
import cv2 as cv
import time
import serial


class CustomEnv(gym.Env):
    """Custom Environment that follows gym interface."""

    metadata = {"render_modes": ["human"], "render_fps": 30}

    def __init__(self, tty_usb='/dev/ttyUSB2', discount_f=7):
        super().__init__()
        # Define action and observation space
        # They must be gym.spaces objects
        # Example when using discrete actions:
        self.action_space = spaces.Box(low=-1, high=1, shape=(1,), dtype=np.float32)
        # Example for using image as input (channel-first; channel-last also works):
        self.observation_space = spaces.Box(low=0, high=255,
                                            shape=(1, 480, 640), dtype=np.uint8)

        self._latest_obs = np.ndarray((1, 480, 640),dtype=np.uint8)
        self._latest_img = None
        self._cur_action = None
        self._last_action = None

        self._max_vel = np.pi
        self._ser = serial.Serial(tty_usb, 115200)
        self._discount_factor = discount_f

    def forward_img(self, img):
        self._latest_img = img
        self._latest_obs = img[np.newaxis, : , :]
        return
    
    def cal_rwd(self):
        if self._latest_img is None:
            return 0.0

        img_rgb = cv.cvtColor(self._latest_img,cv.COLOR_GRAY2RGB)
        orb = cv.ORB.create(nfeatures=1000,fastThreshold=20)
        kp = orb.detect(img_rgb)

        num_kp = len(kp)
        X = np.array([[p.pt[0],p.pt[1]] for p in kp])
        sig_num = -1.0 if len(kp)<self._feature_nums else 1.0

        count_dis = 1
        avg_cor = np.array([0.0,0.0])

        for i in range(num_kp):
            p = X[i]
            cv.circle(img_rgb,(int(p[0]),int(p[1])),3,(0,255,0),-1)
            avg_cor = avg_cor + (p-avg_cor)/count_dis
            count_dis += 1

        img_rgb = cv.resize(img_rgb,(320,240))
        cv.imshow("reward_img",img_rgb)
        cv.waitKey(1)

        num_weight = (self._feature_nums-num_kp)*(self._feature_nums+num_kp)/(self._feature_nums)**2
    

        u = self._latest_img.shape[1]/2
        dis_weight = avg_cor[0]*(2*u-avg_cor[0])/u**2

        vel_now = self._max_vel * self._cur_action
        vel_last = self._last_action * self._max_vel
        vel_dif = abs(vel_now - vel_last)

        power_weight = 0.0
        #power_weight -= 1.0 * (vel_now/np.pi)**2
        power_weight -= 1.0 * (0.5 * vel_dif/np.pi)**2

        reward = sig_num * (dis_weight if sig_num>0 else num_weight) + power_weight
        return reward

    def step(self, action):
        self._cur_action = action[0]
        reward = self.cal_rwd()
        self._last_action = self._cur_action

        # Serial command calculation
        speed = max(min(action[0],1),-1)
        speed = speed* (32767) / self._discount_factor  # 速度，单位是rad
        speed_high = (int(speed ) >> 8) & 0xFF  # 取高八位
        speed_low = int(speed ) & 0xFF  # 取低八位
        data = [0xAA, 0xAF, 0x00, 0x02, 0x02, speed_high, speed_low]
        checksum = sum(data) & 0xFF
        data.append(checksum)
        data_bytes = bytes(data)
        # self._ser.write(data_bytes)

        return self._latest_obs, reward, False, False, {}

    def reset(self, seed=None, options=None):
        self.seed = seed
        return self._latest_obs, {}

    def render(self):
        return True

    def close(self):
        return True